Rajesh3338 commited on
Commit
8c6bfb1
·
verified ·
1 Parent(s): 86930a3

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +19 -84
app.py CHANGED
@@ -1,87 +1,22 @@
1
  import gradio as gr
2
- import torch
3
- from langchain_huggingface import HuggingFaceEmbeddings
4
- from langchain_community.document_loaders import TextLoader
5
- from langchain_community.vectorstores import FAISS
6
- from langchain.text_splitter import RecursiveCharacterTextSplitter
7
- from langchain.chains import RetrievalQA
8
- from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
9
- from langchain_huggingface import HuggingFacePipeline
10
  import spaces
11
- zero = torch.Tensor([0]).cuda()
12
- print(zero.device) # This will likely print 'cpu'
13
- device = "cuda" if torch.cuda.is_available() else "cpu"
14
-
15
- # Load and process the document
16
- doc_loader = TextLoader("dataset.txt")
17
- docs = doc_loader.load()
18
- text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
19
- split_docs = text_splitter.split_documents(docs)
20
-
21
- # Create embeddings and vector store
22
- embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
23
- vectordb = FAISS.from_documents(split_docs, embeddings)
24
-
25
- # Load model and tokenizer
26
- model_name = "01-ai/Yi-Coder-9B-Chat"
27
- tokenizer = AutoTokenizer.from_pretrained(model_name)
28
- model = AutoModelForCausalLM.from_pretrained(
29
- model_name,
30
- device_map="auto",
31
- torch_dtype=torch.float16 if device == "cuda" else torch.float32
32
- )
33
-
34
- # Set up the QA pipeline
35
- qa_pipeline = pipeline(
36
- "text-generation",
37
- model=model,
38
- tokenizer=tokenizer,
39
- max_new_tokens=750,
40
- pad_token_id=tokenizer.eos_token_id
41
- )
42
-
43
- llm = HuggingFacePipeline(pipeline=qa_pipeline)
44
-
45
- # Set up the retriever and QA chain
46
- retriever = vectordb.as_retriever(search_kwargs={"k": 5})
47
- qa_chain = RetrievalQA.from_chain_type(
48
- retriever=retriever,
49
- chain_type="stuff",
50
- llm=llm,
51
- return_source_documents=False
52
- )
53
-
54
- def preprocess_query(query):
55
- if "script" in query or "code" in query.lower():
56
- return f"Write a CPSL script: {query}"
57
- return query
58
-
59
- def clean_response(response):
60
- result = response.get("result", "")
61
- if "Answer:" in result:
62
- return result.split("Answer:")[1].strip()
63
- return result.strip()
64
-
65
- def chatbot_response(user_input):
66
- processed_query = preprocess_query(user_input)
67
- raw_response = qa_chain.invoke({"query": processed_query})
68
- return clean_response(raw_response)
69
-
70
- # Gradio interface
71
- with gr.Blocks() as chat_interface:
72
- gr.Markdown("# CPSL Chatbot")
73
- chat_history = gr.Chatbot(type='messages')
74
- user_input = gr.Textbox(label="Your Message:")
75
- send_button = gr.Button("Send")
76
-
77
- def interact(user_message, history):
78
- bot_reply = chatbot_response(user_message)
79
- history.append({"role": "user", "content": user_message})
80
- history.append({"role": "assistant", "content": bot_reply})
81
- return history, history
82
-
83
- send_button.click(interact, inputs=[user_input, chat_history], outputs=[chat_history, chat_history])
84
 
85
- # Launch the interface
86
- if __name__ == "__main__":
87
- chat_interface.launch()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  import gradio as gr
 
 
 
 
 
 
 
 
2
  import spaces
3
+ import torch
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4
 
5
+ def gpu_info():
6
+ info = []
7
+ info.append(f"Torch Version: {torch.__version__}")
8
+ info.append(f"CUDA Available: {torch.cuda.is_available()}")
9
+ info.append(f"GPU Count: {torch.cuda.device_count()}")
10
+
11
+ if torch.cuda.is_available():
12
+ info.append(f"GPU Name: {torch.cuda.get_device_name(0)}")
13
+ info.append(f"VRAM: {torch.cuda.get_device_properties(0).total_memory // (1024**2)} MB")
14
+ else:
15
+ info.append("No GPU detected.")
16
+
17
+ return "\n".join(info)
18
+
19
+ # Gradio Interface
20
+ iface = gr.Interface(fn=gpu_info, inputs=None, outputs="text", title="GPU Information Checker")
21
+
22
+ iface.launch()